#include "OptimiseRt.h"
#include "levmar.h"


double computeInitalDamping(const Mat&JtJ)
{
	double t = 1;
	double max = -1;
	for (size_t i = 0; i < JtJ.rows; i++)
	{
		double tmp = JtJ.ptr<double>(i)[i];
		if (tmp>max)
		{
			max = tmp;
		}
	}
	return max*t;
}

Mat constructG(const Mat&J, const Mat &f)
{
	return J.t()*f;
}
Mat constructA(const Mat &J)
{
    Mat JtJ;
    Mat Jt;
    Jt = J.t();
    JtJ = Mat::zeros(J.cols,J.cols,J.type());
    for (int i=0; i<J.cols; ++i)
    {
        JtJ += Jt.col(i)*J.row(i);
    }
    return JtJ;
}
bool isConvergence(const Mat &f,const double epsilon = 1.0)
{
	return normInf<double, double>((double*)f.data, f.rows)<epsilon;
}

Mat computeStep(const Mat&A,const double mu,const Mat &G)
{
	Mat a;
	a = A + Mat::eye(A.rows, A.cols, CV_64F)*mu;
	return -1 * a.inv()*G;
}


void updateStep(const Mat&delta, double pp[8])
{
	double * p = (double*)delta.data;
	for (size_t i = 0; i < 8; i++)
	{
		pp[i] += p[i];
	}
}

void updateRt(const double p[8],Matx33d &R,Matx31d &t)
{
    Matx33d Rx = Matx33d::zeros();
    auto d = (double*)Rx.val;
    d[1] = -p[2];
    d[3] = p[2];
    d[2] = p[1];
    d[6] = -p[1];
    d[5] = -p[0];
    d[7] = p[0];
    R = R-Rx*R;
    t(0)+=p[3];
    t(1)+=p[4];
    t(2)+=p[5];
}

Matx33d skewSymmetricMatrix(double a, double b, double c)
{
    Matx33d Rx = Matx33d::zeros();
	double * p = Rx.val;
	p[1] = -c;
	p[3] = c;
	p[2] = b;
	p[6] = -b;
	p[5] = -a;
	p[7] = a;
	return Rx;
}

void verifyJacobian(const vector<MarkerPixel> &M, const Mat &I, IntrinsicMatrix&K, Camera&C,const double * p, const vector<Point3f> &kp3, const vector<Point2f> &kp2,double w1, double w2,double scale,const Mat &J)
{
    vector<double> err1,err2;
    Mat f = computeEnergyLM(M, I, K, C, p, kp3, kp2, w1, w2,scale);
    vector<double> err11,err22;
	double step = 0.01;
	Matx33d dR = skewSymmetricMatrix(0,0.01,0);
    Camera C2 = C;
	C2.R = C.R - dR*C.R;
//	C2.t.at<double>(2, 0) = C2.t.at<double>(2, 0)+0.05;
//	p[1] += 0.01;
	Mat f2=computeEnergyLM(M, I, K, C2, p, kp3, kp2, w1, w2,scale);
    vector<double> j;
	double *ff = (double*)f.data;
	double *ff2 = (double*)f2.data;
    for (int i=0; i<J.rows; ++i) {
		cout<<(ff2[i] - ff[i]) / step;
			cout<<"\t"<<J.at<double>(i,1)<<endl;
     } 
}

Mat computeEnergyLM(const vector<MarkerPixel> &M, const Mat &I, IntrinsicMatrix&K, Camera&C,const double *p, const vector<Point3f> &kp3, const vector<Point2f> &kp2, double w1, double w2,double scale)
{
	Camera C2 = C;
	C2.applyTranform(p);
	vector<Point3f> kp3_c;
	C2.TransformToCameraPlane(kp3, kp3_c);
	vector<Point2f> kp2_d;
	K.projectToImage(kp3_c, kp2_d);
	double e = 0, e2 = 0;
	Mat f(kp2_d.size() * 2 + M.size(), 1, CV_64F);
	double *err1 = (double*)f.data;
	double *err2 = err1 + M.size();
	for (size_t i = 0; i < kp2_d.size(); i++)
	{
		err2[2 * i] = w2*(kp2_d[i].x - kp2[i].x);
		err2[2 * i+1] = w2*(kp2_d[i].y - kp2[i].y);
	}

	const int m_x = I.cols;
	const int m_y = I.rows;
	Point3f tt;
	Point2i ttt;
	IntrinsicMatrix K2 = K.applyScale(scale);
	for (size_t i = 0; i < M.size(); i++)
	{
		auto m = M[i];
		C2.TransformToCameraPlane(m.m_3d, tt);
		K2.projectToImage(tt, ttt);
		
		if (i == 1)
		{
			cout << ttt.x << " " << ttt.y << endl;
		}
		if (ttt.x<0 || ttt.x >= m_x || ttt.y<0 || ttt.y >= m_y)
		{
			err1[i] = -1;
			continue;
		}
		err1[i] = w1*(p[6] * I.ptr<uchar>(ttt.y)[ttt.x] + p[7] - m.i);
	}
	return f;
}

void computeJacobian(const vector<MarkerPixel> &M, const Mat &I,const Mat&gx,const Mat&gy, IntrinsicMatrix&K, Camera &C, const double*p, const vector<Point3f> &kp3, const vector<Point2f> &kp2, double w1, double w2, double scale,Mat &J)
{
	Camera C2 = C;
	C2.applyTranform(p);
	Point3f c;
	Point2i p_i;
	const auto &tt = C.t.val;
	double rx, ry, rz;
	int cnt = M.size()+kp2.size()*2;
	int cnt1 = M.size();
	J = Mat(cnt, 8, CV_64F, 0.0);
	double e0;
	double ggx, ggy;
	IntrinsicMatrix K2 = K.applyScale(scale);
	for (size_t i = 0; i < M.size(); i++)
	{
		auto &m = M[i];
		C2.TransformToCameraPlane(m.m_3d, c);
		K2.projectToImage(c, p_i);
		auto j2 = (double*)J.ptr<double>(i);
		if (p_i.x<0 || p_i.x >= I.cols || p_i.y<0 || p_i.y >= I.rows)
			continue;
		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		ggx = gx.ptr<float>(p_i.y)[p_i.x];
		ggy = gy.ptr<float>(p_i.y)[p_i.x];
		j2[3] = w1*(p[6] * K2.fx * ggx) / c.z;
		j2[4] = w1*(p[6] * K2.fy * ggy) / c.z;
		j2[5] = w1*((p[6] * ggx) *(K2.cx  - p_i.x) / c.z +(p[6] * ggy)*(K2.cy - p_i.y) / c.z);
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = I.ptr<uchar>(p_i.y)[p_i.x];
		j2[7] = 1;
	}
	vector<Point3f> kp3_c;
	C2.TransformToCameraPlane(kp3, kp3_c);
	vector<Point2f> kp2_d;
	K.projectToImage(kp3_c, kp2_d);

	Point2f dx;
	for (int i = 0; i < kp2_d.size(); i++)
	{
		const auto &t = kp2_d[i];
		dx.x = t.x - kp2[i].x;
		dx.y = t.y - kp2[i].y;
		const auto &c = kp3_c[i];
		auto j2 = (double*)J.ptr<double>(2*i + cnt1);
		j2[3] = w2*K.fx / c.z;
		j2[4] = 0;
		j2[5] = w2*(K.cx - t.x) / c.z;
		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = 0;
		j2[7] = 0;

		j2 = (double*)J.ptr<double>(2 * i + 1 + cnt1);
		j2[3] = 0;
		j2[4] = w2*(K.fy) / c.z;
		j2[5] = w2*(K.cy - t.y) / c.z;

		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = 0;
		j2[7] = 0;
	}	

}

Mat functionF(const Mat&f)
{
	return 1 / 2.0*f.t()*f;
}
Mat functionl(const Mat&delta, const Mat &J, const Mat &f)
{
	return f + J*delta;
}
Mat functionL(const Mat&l)
{
	return 1 / 2.0*l.t()*l;
}

//void computeEnergyLM(const vector<MarkerPixel> &M, const Mat &I, IntrinsicMatrix&K, Camera&C, const float * p, const vector<Point3f> &kp3, const vector<Point2f> &kp2, vector<double> &err1,vector<double> &err2,double w1, double w2)
//{
//	vector<Point3f> kp3_c;
//	C.TransformToCameraPlane(kp3, kp3_c);
//	vector<Point2f> kp2_d;
//	K.projectToImage(kp3_c, kp2_d);
//	double e = 0, e2 = 0;
//    err2.resize(kp2_d.size());
//	for (size_t i = 0; i < kp2_d.size(); i++)
//	{
//		err2[i]=w2*normL2Sqr(kp2_d[i], kp2[i]);
//	}
//
//	const int m_x = I.cols;
//	const int m_y = I.rows;
//	Point3f tt;
//	Point2i ttt;
//    err1.resize(M.size());
//
//	for (size_t i = 0; i < M.size(); i++)
//	{
//		auto m = M[i];
//		C.TransformToCameraPlane(m.m_3d, tt);
//		K.projectToImage(tt, ttt);
//        if(ttt.x<0||ttt.x>=m_x||ttt.y<0||ttt.y>=m_y )
//        {
//            err1[i]=-1;
//            continue;
//        }
//        err1[i] = w1*pow(p[0] * I.ptr<uchar>(ttt.y)[ttt.x] + p[1] - m.i, 2);
//	}
//}

void computeGradient(const Mat &_src, Mat &gx, Mat &gy)
{
	Mat src;
	if (_src.type() == CV_8UC3)
	{
		cvtColor(_src, src, CV_BGR2GRAY);
	}
	else
	{
		src = _src.clone();
	}
	gx.create(src.size(), CV_32F);
	gy.create(src.size(), CV_32F);
	for (int i = 1; i < src.rows - 1; i++)
	{
		auto p0 = src.ptr<uchar>(i - 1);
//		auto p1 = src.ptr<uchar>(i);
		auto p2 = src.ptr<uchar>(i + 1);
		auto g = gy.ptr<float>(i);
		for (size_t j = 0; j < src.cols; j++)
		{
			g[j] = (p2[j] - p0[j]) / 2.0;
		}
	}
	auto g = gy.ptr<float>(0);
	auto g2 = gy.ptr<float>(gy.rows - 1);
	auto p0 = src.ptr<uchar>(0);
	auto p1 = src.ptr<uchar>(1);
	auto p2 = src.ptr<uchar>(src.rows - 2);
	auto p3 = src.ptr<uchar>(src.rows - 1);
	for (size_t j = 0; j < src.cols; j++)
	{
		g[j] = (p1[j] - p0[j]) / 2.0;
		g2[j] = (p3[j] - p2[j]) / 2.0;
	}

	for (int i = 0; i < src.rows; i++)
	{
		auto p0 = src.ptr<uchar>(i);
		auto g = gx.ptr<float>(i);
		for (size_t j = 1; j < src.cols - 1; j++)
		{
			g[j] = (p0[j + 1] - p0[j - 1]) / 2.0;
		}
	}

	for (int i = 0; i < src.rows; i++)
	{
		const int cols = gx.cols;
		auto g = gx.ptr<float>(i);
		auto p = gx.ptr<uchar>(i);
		g[0] = (p[1] - p[0]) / 2;
		g[cols - 1] = (p[cols - 1] - p[cols - 2]) / 2.0;
	}
}

//void computeJacobian(const vector<MarkerPixel> &M, const Mat &I, IntrinsicMatrix&K, Camera &C, const float*p, const vector<Point3f> &kp3, const vector<Point2f> &kp2,const vector<double> &err1,const vector<double> &err2, double w1, double w2, Mat &J)
//{
//	Point3f c;
//	Point2i p_i;
//	const auto &tt = (double*)C.t.data;
//	double rx, ry, rz;
//    int cnt = err1.size()+err2.size();
//    int cnt1 = err1.size();
//	J = Mat(cnt, 8,CV_64F, 0.0);
//	double e0;
//	for (size_t i = 0; i < M.size(); i++)
//	{
//		auto &m = M[i];
//		C.TransformToCameraPlane(m.m_3d, c);
//		K.projectToImage(c, p_i);
//        auto j2 = (double*)J.ptr<double>(i);
//        if(p_i.x<0||p_i.x>=I.cols||p_i.y<0||p_i.y>=I.rows)
//            continue;
//		e0 = p[0] * I.ptr<uchar>(p_i.y)[p_i.x] + p[1] - m.i;
//		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
//		j2[3] = 2 * e0*w1*(p[0] * K.fx * m.gx) / c.z;
//		j2[4] = 2 * e0*w1*(p[0] * K.fy * m.gy) / c.z;
//		j2[5] = 2 * e0*w1*(-(p[0]*m.gx) *(K.cx-p_i.x) / c.z - (p[0]*m.gy)*(K.cy-p_i.y)/ c.z);
//		j2[0] = j2[4] * rz - j2[5] * ry;
//		j2[1] = j2[5] * rx - j2[3] * rz;
//		j2[2] = j2[3] * ry - j2[4] * rx;
//        j2[6] = I.ptr<uchar>(p_i.y)[p_i.x];
//        j2[7] = 1;
//	}
//	vector<Point3f> kp3_c;
//	C.TransformToCameraPlane(kp3, kp3_c);
//	vector<Point2f> kp2_d;
//	K.projectToImage(kp3_c, kp2_d);
//    
//    Point2f dx;
//	for (int i = 0; i < kp2_d.size(); i++)
//	{
//		const auto &t = kp2_d[i];
//        dx.x = t.x - kp2[i].x;
//        dx.y = t.y - kp2[i].y;
//        const auto &c = kp3_c[i];
//        auto j2 = (double*)J.ptr<double>(i+cnt1);
//        j2[3] = 2 * w2*dx.x*K.fx/c.z;
//		j2[4] = 2 * w2*(K.fy*dx.y) / c.z;
//		j2[5] = 2 * w2*dx.x*(K.cx-t.x)/c.z+2*w2*dx.y*(K.cy-t.y)/c.z;
//        
//        rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
//		j2[0] = j2[4] * rz - j2[5] * ry;
//		j2[1] = j2[5] * rx - j2[3] * rz;
//		j2[2] = j2[3] * ry - j2[4] * rx;
//        j2[6]=0;
//        j2[7]=0;
//	}
//
//}

Camera::Camera(const Camera & b)
{
	this->R = b.R;
	this->t = b.t;
}

void Camera::applyTranform(const double * p)
{
    Matx33d Rx = Matx33d::zeros();
	auto d = Rx.val;
	d[1] = -p[2];
	d[3] = p[2];
	d[2] = p[1];
	d[6] = -p[1];
	d[5] = -p[0];
	d[7] = p[0];
	R = R - Rx*R;
	t(0) += p[3];
	t(1) += p[4];
	t(2) += p[5];
}

IntrinsicMatrix::IntrinsicMatrix(const IntrinsicMatrix & b)
{
	K = b.K;
	fx = b.fx;
	cx = b.cx;
	fy = b.fy;
	cy = b.cy;
}

IntrinsicMatrix IntrinsicMatrix::applyScale(double scale)
{
	Matx33d k = K*scale;
	k(8) = 1;
	return IntrinsicMatrix(k);
}


void test(const Point3f&s, const IntrinsicMatrix &K, const Camera&C,Point2i &rr)
{
	const auto &r = C.R.val;
	const auto &tt = C.t.val;
	Point3f d;
	d.x = s.x*r[0] + s.y*r[1] + s.z*r[2] + tt[0];
	d.y = s.x*r[3] + s.y*r[4] + s.z*r[5] + tt[1];
	d.z = s.x*r[6] + s.y*r[7] + s.z*r[8] + tt[2];
	rr.x = d.x * K.fx / d.z + K.cx;
	rr.y = d.y * K.fy / d.z + K.cy;
}

void _func(double *p, double *hx, int m, int n, void *adata)
{
    static TickMeter tm2,tm;
    static TickMeter tm3;
    tm.start();
    auto pointers = (char **)adata;
    auto M = (vector<MarkerPixel>*)pointers[0];
    auto I = (Mat*)pointers[1];
    auto K = (IntrinsicMatrix*)pointers[2];
    auto C = (Camera*)pointers[3];
    auto p3 = (vector<Point3f>*)pointers[4];
    auto p2 = (vector<Point2f>*)pointers[5];
    auto w = (double*)pointers[6];

	auto scale = w[2];
	Camera C2 = *C;
	C2.applyTranform(p);
	static vector<Point3f> kp3_c;
	kp3_c.clear();
	C2.TransformToCameraPlane(*p3, kp3_c);
	static vector<Point2f> kp2_d;
	kp2_d.clear();
	K->projectToImage(kp3_c, kp2_d);
	double *err1 = hx;
	double *err2 = err1 + M->size();
	for (size_t i = 0; i < kp2_d.size(); i++)
	{
		err2[2 * i] = w[1]*(kp2_d[i].x - (*p2)[i].x);
		err2[2 * i + 1] = w[1]*(kp2_d[i].y - (*p2)[i].y);
	}

	const int m_x = I->cols;
	const int m_y = I->rows;
	Point3f tt;
	IntrinsicMatrix K2 = K->applyScale(scale);
	tm2.start();
	static vector<Point2i> ttt;
	ttt.resize(M->size());
	for (size_t i = 0; i < M->size(); i++)
	{
		auto &m = (*M)[i];
		//C2.TransformToCameraPlane(m.m_3d, tt);
		//K2.projectToImage(tt, ttt[i]);
		test(m.m_3d, K2, C2, ttt[i]);
	}
	tm3.start();
	for (size_t i = 0; i < M->size(); i++)
	{
		if (ttt[i].x<0 || ttt[i].x >= m_x || ttt[i].y<0 || ttt[i].y >= m_y)
		{
			err1[i] = -1;
			continue;
		}
		err1[i] = w[0]*(p[6] * I->ptr<uchar>(ttt[i].y)[ttt[i].x] + p[7] - (*M)[i].i);
	}
	tm2.stop();
	tm.stop();
	tm3.stop();
	//cout << tm.getTimeSec() << endl;
	//cout << tm2.getTimeSec() << endl;
	//cout << tm3.getTimeSec() << endl;

}

void _jacf(double *p, double *j, int mm, int n, void *adata)
{
	static TickMeter tm;
	tm.start();
	auto pointers = (char **)adata;
	auto M = (vector<MarkerPixel>*)pointers[0];
	auto I = (Mat*)pointers[1];
	auto K = (IntrinsicMatrix*)pointers[2];
	auto C = (Camera*)pointers[3];
	auto p3 = (vector<Point3f>*)pointers[4];
	auto p2 = (vector<Point2f>*)pointers[5];
	auto w = (double*)pointers[6];
	auto gx = (Mat*)pointers[7];
	auto gy = (Mat*)pointers[8];
	auto scale = w[2];

	Camera C2 = *C;
	C2.applyTranform(p);
	Point3f c;
	Point2i p_i;
	const auto &tt = C->t.val;
	double rx, ry, rz;

	double ggx, ggy;
	IntrinsicMatrix K2 = K->applyScale(scale);
	for (size_t i = 0; i < M->size(); i++)
	{
		auto &m = (*M)[i];
		C2.TransformToCameraPlane(m.m_3d, c);
		K2.projectToImage(c, p_i);
		auto j2 = j + mm*i;
		if (p_i.x < 0 || p_i.x >= I->cols || p_i.y < 0 || p_i.y >= I->rows)
		{
			for (size_t j = 0; j < 8; j++)
			{
				j2[j] = 0;
			}
			continue;
		}
		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		ggx = gx->ptr<float>(p_i.y)[p_i.x];
		ggy = gy->ptr<float>(p_i.y)[p_i.x];
		j2[3] = w[0]*(p[6] * K2.fx * ggx) / c.z;
		j2[4] = w[0]*(p[6] * K2.fy * ggy) / c.z;
		j2[5] = w[0]*((p[6] * ggx) *(K2.cx - p_i.x) / c.z + (p[6] * ggy)*(K2.cy - p_i.y) / c.z);
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = I->ptr<uchar>(p_i.y)[p_i.x];
		j2[7] = 1;
	}
	static vector<Point3f> kp3_c;
	kp3_c.clear();
	C2.TransformToCameraPlane(*p3, kp3_c);
	static vector<Point2f> kp2_d;
	kp2_d.clear();
	K->projectToImage(kp3_c, kp2_d);

	Point2f dx;
    int cnt = M->size();
	for (int i = 0; i < kp2_d.size(); i++)
	{
		const auto &t = kp2_d[i];
		dx.x = t.x - (*p2)[i].x;
		dx.y = t.y - (*p2)[i].y;
		const auto &c = kp3_c[i];
		auto j2 = j + mm*cnt + 2*i*mm;
		j2[3] = w[1]*K->fx / c.z;
		j2[4] = 0;
		j2[5] = w[1]*(K->cx - t.x) / c.z;
		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = 0;
		j2[7] = 0;

		j2 = j + mm*cnt + (2 * i+1)*mm;
		j2[3] = 0;
		j2[4] = w[1]*(K->fy) / c.z;
		j2[5] = w[1]*(K->cy - t.y) / c.z;

		rx = c.x - tt[0]; ry = c.y - tt[1]; rz = c.z - tt[2];
		j2[0] = j2[4] * rz - j2[5] * ry;
		j2[1] = j2[5] * rx - j2[3] * rz;
		j2[2] = j2[3] * ry - j2[4] * rx;
		j2[6] = 0;
		j2[7] = 0;
	}

	tm.stop();
//	cout << tm.getTimeSec() << endl;
}

void OptimiseRt::optimise(const cv::Mat &nowImg, const IntrinsicMatrix &K, Camera &C, const vector<Point3f> &pts3d, const vector<Point2f> &pts2d, const cv::Mat &gx, const cv::Mat &gy, double w[3])const
{
    char * tt[9];
    tt[0] = (char*)&M;
    tt[1] = (char*)&nowImg;
    tt[2] = (char*)&K;
    tt[3] = (char*)&C;
    tt[4] = (char*)&pts3d;
    tt[5] = (char*)&pts2d;
    tt[6] = (char*)&w[0];
    tt[7] = (char*)&gx;
    tt[8] = (char*)&gy;

    int n = M.size() + pts3d.size() * 2;
    int mm = 8;
    double info[9];
    double opts[4] = { 0.1,1,0.008,100 };
    double p[8] = {0,0,0,0,0,0,1,0};
    int haha = dlevmar_der(_func,_jacf, p, NULL, mm, n, 100, opts, info, NULL, NULL, (void*)tt);
    cout<<"Iteratration number"<<haha<<endl;
    updateRt(p,C.R,C.t);

}
